import torch

from api import torch_build as tb
from ultralytics import YOLO
from utils.config import MODEL_PATH,PH_MODEL_PATH,H2O2_MODEL_PATH,LE_MODEL_PATH,PIP_MODEL_PATH,SNA_MODEL_PATH,OA_MODEL_PATH,NAG_MODEL_PATH


class TorchPredict(object):
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.load_model()
    def load_model(self):
        self.model = YOLO(MODEL_PATH)
        self.classify_models = {}
        self.classify_models['PH'] = YOLO(PH_MODEL_PATH)
        self.classify_models['H2O2'] = YOLO(H2O2_MODEL_PATH)
        self.classify_models['LE'] = YOLO(LE_MODEL_PATH)
        self.classify_models['PIP'] = YOLO(PIP_MODEL_PATH)
        self.classify_models['SNA'] = YOLO(SNA_MODEL_PATH)
        self.classify_models['OA'] = YOLO(OA_MODEL_PATH)
        self.classify_models['NAG'] = YOLO(NAG_MODEL_PATH)




    def predict(self, request_body_dict: dict):
        data = tb.process(request_body_dict)
        return tb.predict(data, self.model,self.classify_models)

